import os
import shutil

import cv2
import numpy as np
from tqdm import tqdm

# 将src_dir_img下的与dst_dir下的同名的文件复制到dst_dir_img下
def copy_files_by_name(src_dir, ref_dir, dst_dir):
    """
    带进度条的文件名匹配复制（忽略后缀）

    参数:
        src_dir: 源目录路径
        ref_dir: 参考目录路径（用于匹配文件名）
        dst_dir: 目标目录路径
    """
    # 获取参考文件名集合（不含后缀）
    ref_files = {os.path.splitext(f)[0] for f in os.listdir(ref_dir)
                 if os.path.isfile(os.path.join(ref_dir, f))}

    # 创建目标目录
    os.makedirs(dst_dir, exist_ok=True)

    # 获取待处理文件列表
    src_files = [f for f in os.listdir(src_dir)
                 if os.path.isfile(os.path.join(src_dir, f))]

    # 初始化进度条
    with tqdm(src_files, desc="过滤标签对应的原图小图", unit="file", ncols=100) as pbar:
        for filename in pbar:
            src_path = os.path.join(src_dir, filename)
            name_only = os.path.splitext(filename)[0]

            if name_only in ref_files:
                dst_path = os.path.join(dst_dir, filename)
                shutil.copy(src_path, dst_path)
                pbar.set_postfix({'当前文件': filename[:15] + '...'})


# 使用示例
# copy_files_by_name(
#     src_dir="path/to/src_dir_img",
#     ref_dir="path/to/dst_dir",
#     dst_dir="path/to/dst_dir_img"
# )


# 使用示例（自动创建目标目录）
# os.makedirs("目标目录", exist_ok=True)
# copy_files_with_progress("源目录", "目标目录")


# 获取文件夹下的文件数量
def count_files(folder_path):
    """
    统计文件夹下的文件数量

    参数:
        folder_path (str): 文件夹路径

    返回:
        int: 文件数量
    """
    file_count = len([f for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f))])
    return file_count

# 读取中文路径下的图片
def read_image_chinese(path, flags=cv2.IMREAD_GRAYSCALE):
    """支持中文路径的图片读取函数
    参数:
        path (str): 含中文的图片路径
        flags (int): OpenCV读取模式，默认灰度模式
    返回:
        np.ndarray: 图像矩阵，失败返回None
    """
    try:
        # 二进制方式读取文件
        with open(path, 'rb') as f:
            img_bin = f.read()  # 获取二进制数据
        # 将二进制数据转为numpy数组
        img_array = np.frombuffer(img_bin, np.uint8)
        # 解码图像
        return cv2.imdecode(img_array, flags)
    except Exception as e:
        print(f"读取失败: {e}")
        return None

# 读取灰度图像并判断是否存在白色掩码
def check_white_mask(image_path, threshold=250):
    """
    检测灰度图像中是否存在白色掩码区域（像素值≥threshold）

    参数:
        image_path (str): 图像文件路径
        threshold (int): 白色判定阈值，默认250（接近255的值视为白色）

    返回:
        bool: 存在白色掩码返回True，否则返回False
    """
    # 读取灰度图像（单通道）
    img = read_image_chinese(image_path, cv2.IMREAD_GRAYSCALE)

    # 二值化处理：像素值≥threshold设为255(白)，其余设为0(黑)
    # 使用cv2.THRESH_BINARY阈值类型[7,8](@ref)
    _, mask = cv2.threshold(img, threshold, 255, cv2.THRESH_BINARY)

    # 统计非零像素数（即白色区域面积）[1,4](@ref)
    return cv2.countNonZero(mask) > 0

# 判断图片的白色区域至少是 3*3 是则返回true
def check_white_mask_size(image_path, threshold=250, min_size=3):
    """
    检测灰度图像中是否存在白色掩码区域（像素值≥threshold）且至少是min_size*min_size

    参数:
        image_path (str): 图像文件路径
        threshold (int): 白色判定阈值，默认250（接近255的值视为白色）
        min_size (int): 最小掩码区域大小，默认3（即至少是3*3）

    返回:
       bool: 存在白色掩码且至少是min_size*min
      """
    # 读取灰度图像（单通道）
    img = read_image_chinese(image_path, cv2.IMREAD_GRAYSCALE)
    if check_white_mask(image_path, threshold):
        if cv2.countNonZero(img) >= min_size * min_size:
            return True
        else:
            return False

# # 使用示例
# if check_white_mask("img/1_0_450_73.png"):
#     print("✅ 检测到白色掩码")
# else:
#     print("❌ 未检测到白色掩码")

# 源目录和目标目录
src_dir = "../XJU2_切割256/train_lab_small"
src_dir_img = "../XJU2_切割256/train_img_small"
dst_dir = "train_lab"
dst_dir_img = "train_img"

# 创建目标目录（如果不存在）
os.makedirs(dst_dir, exist_ok=True)
os.makedirs(dst_dir_img, exist_ok=True)

# 获取所有PNG文件并处理
# png_files = [f for f in os.listdir(src_dir) if f.endswith('.png')]
# for filename in tqdm(png_files, desc='正在过滤掉不含补油的标签', unit='file'):
#     src_path = os.path.join(src_dir, filename)
#     if check_white_mask(src_path):
#         shutil.copy(src_path, dst_dir)
print(f"共{count_files(dst_dir)}个标签含有补油数据")

# 过滤标签对应的原图小图
copy_files_by_name(src_dir_img, dst_dir, dst_dir_img)
print(f"共{count_files(dst_dir_img)}个标签对应的原图")

